import itertools
import sys
import warnings
from RandomForest import *
from LoadData import *
warnings.filterwarnings("ignore")

def run_for_Best_Parameter():
    train_x, train_y, val_x, val_y = load_data()
    n_estimators_all = [10, 15, 20, 35, 40, 45, 50]
    max_depth_all = [-1, 10, 30, 50]
    min_samples_split_all = [2, 6]
    min_samples_leaf_all = [3, 5]
    min_split_gain_all = [0.0, 0.2, 0.4, 0.8]
    feature_rate_all = [0.2, 0.4, 0.6, 0.75]
    sample_rate_all = [0.3, 0.4, 0.5, 0.6, 0.7]
    max_r = -100
    
    for (n_estimators, max_depth, min_samples_split, min_samples_leaf, min_split_gain, feature_rate, sample_rate) in \
            itertools.product(n_estimators_all, max_depth_all, min_samples_split_all, min_samples_leaf_all, min_split_gain_all, feature_rate_all, sample_rate_all, '------'):
        print('------parameters: ', n_estimators, max_depth, min_samples_split, min_samples_leaf, min_split_gain, feature_rate, sample_rate)
        rf = RandomForest_Classifier(
                n_estimators=n_estimators, max_depth=max_depth, min_samples_split=min_samples_split,
                min_samples_leaf=min_samples_leaf, min_split_gain=min_split_gain, 
                sample_rate=sample_rate, feature_rate=feature_rate, random_state=7
        )
        rf.fit(train_x, train_y)
        print('-------------------------------')
        print('训练集..............')
        train_pred = rf.predict(train_x)
        train_r = r2_score(train_y, train_pred)
        print('train rate is', train_r)
        print('预测集..............')
        val_pred = rf.predict(val_x)
        val_r = r2_score(val_y, val_pred)
        print('validate rate is', val_r)
        print('-------------------------------')

        if val_r > max_r:
            max_r = val_r
            save_model(rf, 'save/Random_Forest.pickle')
            
    rf = load_model('save/Random_Forest.pickle')
    print('Best param is feature rate', rf.feature_rate, ' item rate', rf.sample_rate)
        # break

def run_for_LocalBest_Parameter():
    feature_rate_all = [0.75]
    sample_rate_all = [0.5]
    train_x, train_y, val_x, val_y = load_data()
    for (sample_rate, feature_rate) in itertools.product(sample_rate_all, feature_rate_all) :
        rf = RandomForest_Classifier(n_estimators=10, max_depth=-1, min_samples_split=2,
            min_samples_leaf=1, min_split_gain=0.0, 
            sample_rate=sample_rate, feature_rate=feature_rate, random_state=2)
        rf.fit(train_x, train_y)
        print('训练集..............')
        print('训练集..............', file=sys.stderr)
        train_pred = rf.predict(train_x)
        train_r = r2_score(train_y, train_pred)
        print(train_r)
        print('预测集..............')
        print('预测集..............', file=sys.stderr)
        val_pred = rf.predict(val_x)
        val_r = r2_score(val_y, val_pred)
        print(val_r)

        if val_r > max_r:
            max_r = val_r
            save_model(rf, 'save/Random_Forest.pickle')
    
    rf = load_model('save/Random_Forest.pickle')
    print('Best param is feature rate', rf.feature_rate, ' item rate', rf.sample_rate)

def run_for_trained_Model():  
    train_x, train_y, val_x, val_y = load_data()
    rf = RandomForest_Classifier(
            n_estimators=20,
            sample_rate=0.7,
            feature_rate=0.5
    )
    print('------parameters: ', rf.n_estimators, rf.max_depth, rf.min_samples_split, rf.min_samples_leaf, rf.min_split_gain, rf.feature_rate, rf.sample_rate, '------')
    rf.fit(train_x, train_y)
    print('训练集..............')
    print('训练集..............', file=sys.stderr)
    train_pred = rf.predict(train_x)
    train_r = r2_score(train_y, train_pred)
    print(train_r)
    print('预测集..............')
    print('预测集..............', file=sys.stderr)
    val_pred = rf.predict(val_x)
    val_r = r2_score(val_y, val_pred)
    print(val_r)
    # for i in rf.trees:
    #     print(i.real_depth)
    save_model(rf, 'save/Random_Forest.pickle')

def test():
    rf = load_model('save/Random_Forest.pickle')
    print(rf.__dict__)
    ori, data = load_test_data()
    prediction = rf.predict(data)
    ori['charges'] = prediction
    ori.to_csv('data/submission.csv', index=False)

if __name__=='__main__':
    # run_for_Best_Parameter()
    # run_for_LocalBest_Parameter()
    run_for_trained_Model()
    test()
